import os
import xml.etree.ElementTree as ET
from collections import defaultdict

import cv2
import torch
from PIL import Image
from tabulate import tabulate
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms


class CustomDataset(Dataset):
    def __init__(self, data_dir, mode="train"):
        self.data_dir = data_dir
        self.image_paths = []
        self.labels = []
        # read the train or val.txt
        with open(os.path.join("splits", mode + ".txt"), "r") as f:
            for line in f:
                img_path = os.path.join(data_dir, line.strip())
                xml_path = os.path.splitext(img_path)[0] + ".xml"
                label = self._parse_xml(xml_path)
                if label is not None:
                    self.image_paths.append(img_path)
                    self.labels.append(label)
                # self.image_paths.append(img_path)

        # for filename in os.listdir(data_dir):
        #     if os.path.splitext(filename)[1] in ['.jpeg','.jpg']:
        # img_path = os.path.join(data_dir, filename)
        # xml_path = os.path.splitext(img_path)[0] + '.xml'
        # label = self._parse_xml(xml_path)
        # if label is not None:
        #     self.image_paths.append(img_path)
        #     self.labels.append(label)
        if mode == "train":
            self.transform = transforms.Compose(
                [
                    transforms.ToPILImage(),
                    transforms.RandomResizedCrop(224, scale=(0.64, 1)),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
                ]
            )
        elif mode == "val":
            self.transform = transforms.Compose(
                [
                    transforms.ToPILImage(),
                    transforms.Resize(256),
                    transforms.CenterCrop(224),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
                ]
            )
        print(f"{len(self.image_paths)} images loaded")
        # Analyze label distribution
        self.label_stats = self._analyze_label_distribution()

        # # Display label distribution
        # self._display_label_distribution()

    def _load_labels(self):
        labels = []
        for image_path in self.image_paths:
            xml_path = os.path.splitext(image_path)[0] + ".xml"
            labels.append(self._parse_xml(xml_path))
        return labels

    def _parse_xml(self, xml_path):
        if not os.path.exists(xml_path):
            return None
        tree = ET.parse(xml_path)
        root = tree.getroot()

        attributes = root.find("object").find("attributes")
        age = attributes.find('attribute[name="age"]').find("value").text
        if age not in [str(i) for i in range(-1, 101)]:
            return None
        else:
            age = int(age)

        gender = attributes.find('attribute[name="gender"]').find("value").text
        if gender not in [str(i) for i in range(-1, 2)]:
            return None
        else:
            gender = int(gender)

        glasses = attributes.find('attribute[name="glasses"]').find("value").text
        if glasses not in [str(i) for i in range(-1, 3)]:
            return None
        else:
            glasses = int(glasses)

        race = attributes.find('attribute[name="race"]').find("value").text
        if race not in [str(i) for i in range(-1, 4)]:
            return None
        else:
            race = int(race)

        emotion = attributes.find('attribute[name="emotion"]').find("value").text
        if emotion not in [str(i) for i in range(-1, 3)]:
            return None
        else:
            emotion = int(emotion)

        mask = attributes.find('attribute[name="mask"]').find("value").text
        if mask not in [str(i) for i in range(-1, 2)]:
            return None
        else:
            mask = int(mask)

        hat = attributes.find('attribute[name="hat"]').find("value").text
        if hat not in [str(i) for i in range(-1, 2)]:
            return None
        else:
            hat = int(hat)

        whiskers = attributes.find('attribute[name="whiskers"]').find("value").text
        if whiskers not in [str(i) for i in range(-1, 2)]:
            return None
        else:
            whiskers = int(whiskers)

        labels = {
            "age": age,
            "gender": gender,
            "glasses": glasses,
            "race": race,
            "emotion": emotion,
            "mask": mask,
            "hat": hat,
            "whiskers": whiskers,
        }

        # because -1 can't be used in loss, +1 for all, but remember to -1 it back when predict
        for key, value in labels.items():
            labels[key] = value + 1
        return labels

    def _analyze_label_distribution(self):
        label_stats = defaultdict(lambda: defaultdict(int))

        for label in self.labels:
            label_stats["age"][label["age"]] += 1
            label_stats["gender"][label["gender"]] += 1
            label_stats["glasses"][label["glasses"]] += 1
            label_stats["race"][label["race"]] += 1
            label_stats["emotion"][label["emotion"]] += 1
            label_stats["mask"][label["mask"]] += 1
            label_stats["hat"][label["hat"]] += 1
            label_stats["whiskers"][label["whiskers"]] += 1

        # Display each category in a separate table
        for category, values in label_stats.items():
            sorted_values = sorted(
                values.items(),
                key=lambda x: (int(x[0]) if str(x[0]).isdigit() else x[0]),
            )
            print(f"\nCategory: {category}")
            print(
                tabulate(
                    sorted_values, headers=["Value", "Count"], tablefmt="fancy_grid"
                )
            )

    # def _display_label_distribution(self):
    #     print("Label Distribution:")
    #     for key, value in self.label_stats:
    #         print(f"{key}: {value}")

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        image = cv2.imread(
            image_path
        )  # You may need to adjust the reading of images based on your needs
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert to RGB if needed
        # image  = Image.open(image_path)
        image = self.transform(image)
        label = self.labels[idx]

        return image, torch.tensor(
            list(label.values())
        )  # Assuming PyTorch format (C, H, W)


if __name__ == "__main__":
    # Example usage:
    data_dir = "data"
    custom_dataset = CustomDataset(data_dir)

    # Access a specific sample
    sample = custom_dataset[0]

    loader = DataLoader(custom_dataset, batch_size=2)
    sample = next(iter(loader))
    print("done")
